from enum import Enum
from functools import cached_property, reduce
from typing_extensions import Self
from pydantic import BaseModel, Field
from pysam import FastaFile  # pylint: disable=E0611
from .seq import Sequence, MTSequence 


class RegionType(str, Enum):
    CDS = "CDS"
    UTR = "UTR"
    INTRON = "Intron"
    UD = "UpDownStream"


class Region(BaseModel):
    chrom: str
    start: int
    end: int
    type: RegionType
    exon: str
    name: str = Field(default="")

    class Config:
        arbitrary_types_allowed = True

    @cached_property
    def size(self) -> int:
        return self.end - self.start + 1

    @classmethod
    def from_transcript(cls, transcript: "Transcript") -> list[Self]:
        tag = 1 if transcript.strand == "+" else -1
        regions = list(
            map(
                lambda i, s, e: Region(chrom=transcript.chrom, start=s + 1, end=e - 1, type=RegionType.INTRON, exon=f"intron{i}", name=f"intron{i}"),
                list(range(1, transcript.exon_count))[::tag],
                transcript.exon_ends[:-1],
                transcript.exon_starts[1:],
            )
        )
        cds, utr = 1, 5
        for i, (start, end) in enumerate(zip(transcript.exon_starts[::tag], transcript.exon_ends[::tag])):
            exon_name, cds_name, utr_name = f"exon{i+1}", f"CDS{cds}", f"UTR{utr}"
            _start, _end = max(start, transcript.cds_start), min(end, transcript.cds_end)
            if _start > _end:
                regions.append(Region(chrom=transcript.chrom, start=start, end=end, type=RegionType.UTR, exon=exon_name, name=utr_name))
            else:
                if _start > start:
                    regions.append(Region(chrom=transcript.chrom, start=start, end=_start - 1, type=RegionType.UTR, exon=exon_name, name=utr_name))
                regions.append(Region(chrom=transcript.chrom, start=_start, end=_end, type=RegionType.CDS, exon=exon_name, name=cds_name))
                cds, utr = cds + 1, 3
                if _end < end:
                    regions.append(Region(chrom=transcript.chrom, start=_end + 1, end=end, type=RegionType.UTR, exon=exon_name, name=utr_name))
        return sorted(regions, key=lambda x: x.start)


class Transcript(BaseModel):
    name: str
    chrom: str
    strand: str
    tx_start: int
    tx_end: int
    cds_start: int
    cds_end: int
    exon_count: int
    exon_starts: list[int]
    exon_ends: list[int]
    gene: str
    sequence: Sequence | MTSequence
    gene_id: str = Field(default="")
    regions: list[Region] = Field(default=[])

    class Config:
        arbitrary_types_allowed = True

    @cached_property
    def pk(self) -> str:
        return f"{self.chrom}:{self.gene}:{self.name}"

    @cached_property
    def is_unk(self) -> bool:
        return self.cds_end - self.cds_start + 1 == 0

    @cached_property
    def has_utr3(self) -> int:
        return (self.cds_end < self.tx_end) if self.strand == "+" else (self.cds_start > self.tx_start)

    @cached_property
    def has_utr5(self) -> int:
        return (self.cds_start > self.tx_start) if self.strand == "+" else (self.cds_end < self.tx_end)

    @cached_property
    def cds_count(self) -> int:
        return len(filter(lambda x: x.type == RegionType.CDS, self.regions))

    @cached_property
    def utr_sizes(self) -> tuple[int, int]:
        utr_size1, utr_size2 = 0, 0
        for region in filter(lambda x: x.type == RegionType.UTR, self.regions):
            if region.end < self.cds_start:
                utr_size1 += region.size
            else:
                utr_size2 += region.size
        return utr_size1, utr_size2

    @cached_property
    def cds_size(self) -> int:
        return sum(map(lambda s: s.size, filter(lambda t: t.type == RegionType.CDS, self.regions)))

    @cached_property
    def dna(self) -> Sequence | MTSequence:
        # SequenceOfDNA Regions中所有Sequence拼接结果即DNA
        return self.sequence

    @cached_property
    def coding_dna(self) -> Sequence | MTSequence:
        # SequenceOfCodingDNA Regions中CDS的Sequence拼接结果即CDNA
        iter_cds = filter(lambda t: t.type == RegionType.CDS, self.regions)
        iter_seq = map(lambda s: self.sequence[s.start - self.tx_start : s.end - self.tx_start + 1], iter_cds)
        return reduce(lambda x, y: x + y, iter_seq)

    def find_region(self, pos: int) -> tuple[Region, int, int]:
        if pos < self.tx_start:
            return (
                Region(
                    chrom=self.chrom,
                    start=self.tx_start - 1,
                    end=self.tx_start - 1,
                    type=RegionType.UD,
                    exon="exon1" if self.strand == "+" else f"exon{self.exon_count}",
                    name=self.regions[0].name,
                ),
                0,
                0,
            )
        if pos > self.tx_end:
            return (
                Region(
                    chrom=self.chrom,
                    start=self.tx_end + 1,
                    end=self.tx_end + 1,
                    type=RegionType.UD,
                    exon=f"exon{self.exon_count}" if self.strand == "+" else "exon1",
                    name=self.regions[-1].name,
                ),
                self.cds_size,
                self.utr_sizes[1],
            )
        cds_size, utr_size = 0, 0
        for region in self.regions:
            if region.start <= pos <= region.end:
                return region, cds_size, utr_size
            if region.type == RegionType.CDS:
                cds_size += region.size
                utr_size = 0
            if region.type == RegionType.UTR:
                utr_size += region.size
        raise ValueError(f"can not determine the region of position: {pos}")

    def with_data(self, genome_fasta: FastaFile, gene_id_dict: dict[str, str]):
        if self.exon_count != len(self.exon_starts) or self.exon_count != len(self.exon_ends):
            raise ValueError("the length of exon starts, exon ends and exon count is not equal")
        if not self.regions:
            self.regions = Region.from_transcript(self)
        if not self.sequence:
            self.sequence = self.sequence.__class__(genome_fasta.fetch(self.chrom, self.tx_start - 1, self.tx_end).upper())
        if not self.gene_id:
            self.gene_id = gene_id_dict and gene_id_dict.get(f"{self.chrom}:{self.gene}", "")

    @classmethod
    def from_genepred(cls, line: str) -> Self:
        row = line.split("\t")
        instance = cls(
            name=row[1],
            chrom=row[2],
            strand=row[3],
            tx_start=int(row[4]) + 1,
            tx_end=int(row[5]),
            cds_start=int(row[6]) + 1,
            cds_end=int(row[7]),
            exon_count=int(row[8]),
            exon_starts=list(map(lambda x: int(x) + 1, row[9].rstrip(",").split(","))),
            exon_ends=list(map(int, row[10].rstrip(",").split(","))),
            gene=row[12],
            sequence=Sequence("") if row[2] not in ["chrM", "MT"] else MTSequence(""),
        )
        return instance

    def calc_utr_distance(self, region: Region, utr_size: int, pos: int) -> str:
        trans_utr_size1, trans_utr_size2 = self.utr_sizes
        if pos < self.cds_start:
            # 编码区上游
            if self.strand == "+":
                # 正链
                if trans_utr_size1 == 0:
                    # 没有UTR区
                    return f"-{self.cds_start-pos}"
                # 有UTR区
                if region.type == RegionType.UTR:
                    # 位置在UTR区
                    return f"-{trans_utr_size1 - (utr_size + pos - region.start)}"
                if region.type == RegionType.INTRON:
                    # 位置在INTRON区
                    dist1, dist2 = pos - region.start + 1, region.end - pos + 1
                    return f"-{trans_utr_size1 - utr_size + 1}+{dist1}" if dist1 < dist2 else f"-{trans_utr_size1 - utr_size}-{dist2}"
                # 位置在UP/DOWN STREAM区
                return f"-{trans_utr_size1}-{self.tx_start-pos}"
            # 负链
            if trans_utr_size1 == 0:
                # 没有UTR区
                return f"*{self.cds_start-pos}"
            # 有UTR区
            if region.type == RegionType.UTR:
                # 位置在UTR区
                return f"*{trans_utr_size1 - (utr_size + pos - region.start)}"
            if region.type == RegionType.INTRON:
                # 位置在INTRON区
                dist1, dist2 = pos - region.start + 1, region.end - pos + 1
                return f"*{trans_utr_size1 - utr_size + 1}-{dist1}" if dist1 < dist2 else f"*{trans_utr_size1 - utr_size}+{dist2}"
            # 位置在UP/DOWN STREAM区
            return f"*{trans_utr_size1}+{self.tx_start-pos}"
        if pos > self.cds_end:
            # 编码区下游
            if self.strand == "+":
                # 正链
                if trans_utr_size2 == 0:
                    # 没有UTR区
                    return f"*{pos-self.cds_end}"
                if region.type == RegionType.UTR:
                    # 位置在UTR区
                    return f"*{utr_size + (pos - region.start + 1)}"
                if region.type == RegionType.INTRON:
                    # 位置在INTRON区
                    dist1, dist2 = pos - region.start + 1, region.end - pos + 1
                    return f"*{utr_size}+{dist1}" if dist1 < dist2 else f"*{utr_size + 1}-{dist2}"
                # 位置在UP/DOWN STREAM区
                return f"*{trans_utr_size2}+{pos-self.tx_end}"
            # 负链
            if trans_utr_size2 == 0:
                # 没有UTR区
                return f"-{pos-self.cds_end}"
            # 有UTR区
            if region.type == RegionType.UTR:
                # 位置在UTR区
                return f"-{utr_size + (pos - region.start + 1)}"
            if region.type == RegionType.INTRON:
                # 位置在INTRON区
                dist1, dist2 = pos - region.start + 1, region.end - pos + 1
                return f"-{utr_size}-{dist1}" if dist1 < dist2 else f"-{utr_size + 1}+{dist2}"
            # 位置在UP/DOWN STREAM区
            return f"-{trans_utr_size2}-{pos-self.tx_end}"
        return ""

    def calc_cds_distance(self, region: Region, cds_size: int, pos: int) -> str:
        if self.cds_start <= pos <= self.cds_end:
            if self.strand == "+":
                # 正链
                if region.type == RegionType.CDS:
                    # 位置在CDS区
                    return f"{cds_size + (pos - region.start + 1)}"
                # 位置在INTRON区
                dist1, dist2 = pos - region.start + 1, region.end - pos + 1
                return f"{cds_size}+{dist1}" if dist1 < dist2 else f"{cds_size+1}-{dist2}"
            # 负链
            if region.type == RegionType.CDS:
                # 位置在CDS区
                return f"{self.cds_size - (cds_size + pos - region.start)}"
            # 位置在INTRON区
            dist1, dist2 = pos - region.start + 1, region.end - pos + 1
            return f"{self.cds_size-cds_size+1}-{dist1}" if dist1 < dist2 else f"{self.cds_size-cds_size}+{dist2}"
        return ""
